import wx, wx.aui
import numpy as np
import matplotlib as mpl
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from dialogs import SizeDialog
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas
from matplotlib.backends.backend_wxagg import NavigationToolbar2WxAgg as Toolbar


class MyToolbar(Toolbar):
    """
    Some tweaks to the default matplotlib toolbar to move and resize
    the canvas scrollbar corresponding to pan/zoom operations
    """
    def __init__(self, panel, canvas):
        Toolbar.__init__(self, canvas)
        self.panel = panel
        self.canvas = canvas
        self.DeleteToolByPos(7)  #Removes the subplot configuration tool
        
    def draw(self):
        Toolbar.draw(self)
        self.panel.updateWindow()
    def dynamic_update(self):
        Toolbar.dynamic_update(self)
        self.panel.updateWindow()
    def save_figure(self, *args, **kwargs):
        oldsize = self.canvas.figure.get_size_inches()
        dlg = SizeDialog(self.panel, self.canvas.figure)
        if dlg.ShowModal() == wx.ID_OK:
            out = dlg.getOutput()
            self.canvas.figure.set_size_inches(out['Width'], out['Height'])
            Toolbar.save_figure(self, *args, **kwargs)
            self.canvas.figure.set_size_inches(oldsize)

        
    

class My_Axes(Axes):
    name = "My_Axes"
    def __init__(self, fig, rect, vprops=None, *args, **kwds):        
        Axes.__init__(self, fig, rect, *args, **kwds)              
        if vprops:
            self.vprops = vprops
        else:
            self.vprops = self.get_xlim()+self.get_ylim()
                        
    def drag_pan(self, button, key, x, y):
        _xmin, _xmax = self.get_xlim()
        _ymin, _ymax = self.get_ylim()
        Axes.drag_pan(self, button, key, x, y)
        
        xmin, xmax = self.get_xlim()
        ymin, ymax = self.get_ylim()
        if xmin < self.vprops[0] or xmax > self.vprops[1]:
            self.set_xlim(_xmin, _xmax)
        if ymin < self.vprops[2] or ymax > self.vprops[3]:
            self.set_ylim(_ymin, _ymax)
        
    def set_win_props(self, vprops):
        self.vprops = vprops     


mpl.projections.register_projection(My_Axes)


class PlotPanel(wx.Panel):
    def __init__(self, parent, Fig=None, vwin=[], *args, **kwds):
        '''Creates a wxPanel with a scrollable figure canvas.
        Fig = a matplotlib figure()
        vwin = the virtual window, (v_xmin, v_xmax, v_ymin, v_ymax) in data coordinates'''
        wx.Panel.__init__(self, parent, *args, **kwds)

        self.makeFig(Fig)

        self.labelAxes()
        self.initProps(vwin)
        if self.vprops:
            self.setScrollbars()
            self.enable_scrolling(True)

        self.statusbar = wx.StatusBar(self, -1)
        self.toolbar = MyToolbar(self, self.canvas)
        #Bindings

        self.canvas.mpl_connect("pick_event", self.OnPick)
        self.canvas.mpl_connect('motion_notify_event', self.OnMove)
        self.canvas.mpl_connect('button_release_event', self.OnRelease)
        #Setup wx sizer
        self.sizer = wx.BoxSizer(wx.VERTICAL)
        self.sizer.Add(self.toolbar, 0, wx.EXPAND | wx.ALL, 5)
        self.sizer.Add(self.canvas, 1, wx.EXPAND | wx.GROW | wx.ALL, 5)
        self.sizer.Add(self.statusbar, 0, wx.ALL | wx.EXPAND, 5)
        self.SetSizer(self.sizer)
        self.Fit()

    def makeFig(self, Fig):
        '''Initializes figure, axes and canvas.  You can override to setup your own figure'''
        if Fig:
            self.Fig = Fig
            self.ax = self.Fig.gca()
            self.canvas = FigureCanvas(self, -1, self.Fig)
        else:
            self.Fig = Figure()
            self.ax = self.Fig.add_subplot(111)
            self.canvas = FigureCanvas(self, -1, self.Fig)


    def initProps(self, vwin):
        '''inits axes and window props from input figure and virtual window'''
        self.axvals = [self.ax.get_xlim()[0], self.ax.get_xlim()[1], 
                       self.ax.get_xticks()[1]-self.ax.get_xticks()[0],
                       self.ax.get_ylim()[0], self.ax.get_ylim()[1], 
                       self.ax.get_yticks()[1]-self.ax.get_yticks()[0] ]
        self.axprops = self.axesDict(self.axvals)
        if vwin:
            self.winvals = vwin
            self.vprops = self.winDict(self.winvals)
        else:
           pass
       
    def labelAxes(self):
        self.ax.set_xlabel("x-axis")
        self.ax.set_ylabel("y-axis")

    def redraw(self):
        '''Updates Plot'''
        self.canvas.draw()

    def axesDict(self, vals):
        '''returns axes ditionary from list of vals'''
        axkeys = ['xmin', 'xmax', 'xticks', 'ymin', 'ymax', 'yticks']
        return dict(zip(axkeys, vals))

    def winDict(self, vwin):
        '''sets window dictionary from vwin list'''
        vkeys = ['xmin', 'xmax', 'ymin', 'ymax']
        return dict(zip(vkeys, vwin))

    def resetAxes(self, props, draw=True):
        self.ax.set_xlim(props['xmin'], props['xmax'])
        self.ax.set_xticks( np.arange(self.vprops['xmin'], self.vprops['xmax'], props['xticks']) )
        self.ax.set_yticks( np.arange(self.vprops['ymin'], self.vprops['ymax'], props['yticks']) )
        self.ax.set_ylim(props['ymin'], props['ymax'])
        if draw: 
            self.canvas.draw()

    def resetWin(self, newprops):
        self.vprops = newprops
        self.winvals = [newprops['xmin'], newprops['xmax'], 
                        newprops['ymin'], newprops['ymax'] ] 
        self.ax.set_win_props(self.winvals)
    
    def updateWindow(self):
        '''Updates the axes property dict corresponding to user navigation
        and positions the scrollbar correctly'''        
        
        self.initProps(None)
        #set scrollbar
        self.setScrollbars()

    def enable_scrolling(self, flag):
        if flag:
            self.canvas.Bind(wx.EVT_SCROLLWIN, self.OnScroll)
            self.canvas.Bind(wx.EVT_MOUSEWHEEL, self.OnMouse)
        else:
            self.canvas.Unbind(wx.EVT_SCROLLWIN)
            self.canvas.Unbind(wx.EVT_MOUSEWHEEL)

    def scroll2data(self, t, orient):
        if orient == wx.VERTICAL:
            return -t + self.vprops['ymax']
        else:
            return t + self.vprops['xmin']

    def data2scroll(self, x , orient):
        if orient == wx.VERTICAL:
            return self.vprops['ymax'] - x
        else:
            return x - self.vprops['xmin']

    def setScrollbars(self):
        self.ywin = self.vprops['ymax'] - self.vprops['ymin']
        self.xwin = self.vprops['xmax'] - self.vprops['xmin']
        self.ypage = self.axprops['ymax'] - self.axprops['ymin']
        self.xpage = self.axprops['xmax'] - self.axprops['xmin']
        xpos = self.data2scroll(self.axprops['xmin'], wx.HORIZONTAL)
        ypos = self.data2scroll(self.axprops['ymax'], wx.VERTICAL)
        self.canvas.SetScrollbar(wx.HORIZONTAL,xpos,self.xpage,self.xwin)
        self.canvas.SetScrollbar(wx.VERTICAL,ypos,self.ypage,self.ywin)

    def OnMouse(self, event):
        '''Mousewheel Event Handler'''
        #Get Scroll Data
        scrollspeed = 10
        amt = event.GetWheelRotation() #Raw data
        units = amt/(-(event.GetWheelDelta())) #Number of scroll units with proper direction
        y = self.canvas.GetScrollPos(wx.VERTICAL) + scrollspeed*units #final position of scrollbar button

        #Update viewing window.
        self.axprops['ymax'] = self.scroll2data(y, wx.VERTICAL)
        self.axprops['ymin'] = self.axprops['ymax'] - self.ypage

        #Enforce absolute min
        if self.axprops['ymin'] < self.vprops['ymin']:
            self.axprops['ymin'] = self.vprops['ymin']
            self.axprops['ymax'] = self.vprops['ymin'] + self.xpage

        #Enforce absolute max
        if self.axprops['ymax'] > self.vprops['ymax']:
            self.axprops['ymax'] = self.vprops['ymax']
            self.axprops['ymin'] = self.vprops['ymax'] - self.ypage

        #Update scrollbar position
        self.canvas.SetScrollPos(wx.VERTICAL, y, True)
        self.resetAxes(self.axprops)

    def OnScroll(self, event):
        '''Event Handler for scrollbar end buttons'''
        t = event.GetPosition()

        if event.GetOrientation() == wx.HORIZONTAL:
            self.axprops['xmin'] = self.scroll2data(t, wx.HORIZONTAL)
            self.axprops['xmax'] = self.axprops['xmin'] + self.xpage
            self.canvas.SetScrollPos(wx.HORIZONTAL, t, True)
        else:
            self.axprops['ymax'] = self.scroll2data(t, wx.VERTICAL)
            self.axprops['ymin'] = self.axprops['ymax'] - self.ypage
            self.canvas.SetScrollPos(wx.VERTICAL, t, True)
        self.resetAxes(self.axprops)

    def OnMove(self, event):
        '''Mouse motion event handler'''
        x = event.xdata
        y = event.ydata
        update = "x = " + str(x)  + ' | y = ' + str(y)
        self.statusbar.SetStatusText(update)

    def OnPick(self, event):
        print "Pick Event" + str(event)

    def OnRelease(self, event):
        print "Release Event: " + str(event)


class PlotNotebook(wx.Panel):
    def __init__(self, parent, id = -1):
        wx.Panel.__init__(self, parent, id=id)
        self.nb = wx.aui.AuiNotebook(self)
        sizer = wx.BoxSizer()
        sizer.Add(self.nb, 1, wx.EXPAND)
        self.SetSizer(sizer)

    def add(self, Fig, win, name="plot"):
       page = PlotPanel(self.nb, Fig, win)
       self.nb.AddPage(page,name)

def demo_fig():
    Fig = Figure()
    ax = Fig.add_subplot(111, projection="My_Axes")

    t = np.arange(-100, 100, .1)
    y = t**2*np.sin(t)
    ax.plot(t,y, picker=5.0)
    ax.set_xlim(-20,20)
    ax.set_ylim(-100,100)
    
   # ax.set_xticks(np.arange(-20, 20, 5))
 #   ax.set_yticks(np.arange(-100, 100, 10))
    win = [-100,100, -10000, 10000]
    ax.set_win_props(win)
    return Fig, win


def demo():
    app = wx.PySimpleApp()
    frame = wx.Frame(None,-1,'PlotPanel')
    plotter = PlotNotebook(frame)
    Fig, win = demo_fig()
    plotter.add(Fig, win)
    frame.SetSize((800,600))

    #axes1.figure.canvas.draw()
    #axes2.figure.canvas.draw()
    frame.Show()
    app.MainLoop()


if __name__ == "__main__": demo()




